#include <type_traits>
#include "kernel_operator.h"
using namespace AscendC;

constexpr int32_t BUFFER_NUM = 2; 

template<typename TYPE_X1> class KernelNonMaxSuppression {
    using T = TYPE_X1;
public:
    __aicore__ inline KernelNonMaxSuppression() {}
    __aicore__ inline void Init(GM_ADDR boxes, GM_ADDR scores, GM_ADDR selected_indices,
            int32_t center_point_box, int32_t num_batches, int32_t num_classes, int32_t spatial_dimension, int32_t max_output_boxes_per_class,
            float iou_threshold, float score_threshold, uint32_t tileDataNum) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->center_point_box = center_point_box;
        this->num_batches = num_batches;
        this->num_classes = num_classes;
        this->spatial_dimension = spatial_dimension;
        this->max_output_boxes_per_class = max_output_boxes_per_class;
        this->iou_threshold = iou_threshold;
        this->score_threshold = score_threshold;
        this->tileDataNum = tileDataNum;
        uint32_t boxeslength = num_batches * spatial_dimension * 4;
        uint32_t boxeslengthaligned = (boxeslength + 7) / 8 * 8;
        uint32_t scoreslength = num_batches * num_classes * spatial_dimension;
        uint32_t scoreslengthaligned = (scoreslength + 7) / 8 * 8;
        uint32_t selectedindiceslength = num_batches * num_classes * 3;
        uint32_t selectedindiceslengthal = (selectedindiceslength + 7) / 8 * 8;
        this->spatial_dimensional = (spatial_dimension + 7) / 8 * 8;
        this->inQboxeslength = (spatial_dimension * 4 + 7) / 8 * 8;

        boxesGm.SetGlobalBuffer((__gm__ T*)boxes, boxeslengthaligned);
        scoresGm.SetGlobalBuffer((__gm__ T*)scores, scoreslengthaligned);
        selected_indicesGm.SetGlobalBuffer((__gm__ int32_t*)selected_indices, selectedindiceslengthal);

        pipe.InitBuffer(inQboxes, BUFFER_NUM, this->inQboxeslength * sizeof(T));
        pipe.InitBuffer(inQscores, BUFFER_NUM, this->spatial_dimensional * sizeof(T));
        //pipe.InitBuffer(outQindices, BUFFER_NUM, spatial_dimensional * sizeof(int32_t));
        pipe.InitBuffer(tmpBufferA1, this->spatial_dimensional * sizeof(float));
        pipe.InitBuffer(tmpBufferA2, this->spatial_dimensional * sizeof(float));
        pipe.InitBuffer(tmpBufferA3, this->spatial_dimensional * sizeof(float));
        pipe.InitBuffer(tmpBufferA4, this->spatial_dimensional * sizeof(float));
        pipe.InitBuffer(tmpBufferA5, this->spatial_dimensional * sizeof(float));

        pipe.InitBuffer(tmpBufferP1, 2 * sizeof(uint32_t));
        pipe.InitBuffer(tmpBufferP2, 2 * sizeof(uint32_t));
        pipe.InitBuffer(tmpBufferP3, 2 * sizeof(uint32_t));
        pipe.InitBuffer(tmpBufferP4, 2 * sizeof(uint32_t));
        pipe.InitBuffer(tmpBufferX1, this->spatial_dimensional * sizeof(float));
        pipe.InitBuffer(tmpBufferX2, this->spatial_dimensional * sizeof(float));
        pipe.InitBuffer(tmpBufferY1, this->spatial_dimensional * sizeof(float));
        pipe.InitBuffer(tmpBufferY2, this->spatial_dimensional * sizeof(float));
        pipe.InitBuffer(tmpBufferSX1, this->spatial_dimensional * sizeof(float));
        pipe.InitBuffer(tmpBufferSX2, this->spatial_dimensional * sizeof(float));
        pipe.InitBuffer(tmpBufferSY1, this->spatial_dimensional * sizeof(float));
        pipe.InitBuffer(tmpBufferSY2, this->spatial_dimensional * sizeof(float));
        pipe.InitBuffer(QueueTmp2, this->spatial_dimensional * sizeof(uint8_t));

        int typeSize = 4; 
        int elementsPerBlock = 32 / typeSize;       // 1个block存放的元素个数
        int elementsPerRepeat = 256 / typeSize;     // 1次repeat可以处理的元素个数

        // 最后确定首次最大repeat值
        int firstMaxRepeat = this->inQboxeslength / 64;           // 此处需要注意：对于tensor高维切分计算接口，firstMaxRepeat就是	repeatTimes；对于tensor前n个数据计算接口，firstMaxRepeat为count/elementsPerRepeat，比如在half类型下firstMaxRepeat就是count/128，在float类型下为count/64，按需填入，对于count<elementsPerRepeat的场景，firstMaxRepeat就是1
        firstMaxRepeat = (firstMaxRepeat < 1) ? 1 : firstMaxRepeat;
        int iter1OutputCount = firstMaxRepeat * 2;                                              // 第一轮操作产生的元素个数
        int iter2AlignStart = RoundUp(iter1OutputCount, elementsPerBlock) * elementsPerBlock; // 第二轮操作起始位置偏移，即第一轮产生的元素个数按照datablock(32字节)向上对齐的结果
        // 第一轮计算完成后，后续可能还需要多轮迭代，此时不可以复用同一块空间，因为第一轮的中间结果索引还需要再进行使用，所以需要继续准备第二轮和第三轮的空间
        int iter2OutputCount = RoundUp(iter1OutputCount, elementsPerRepeat) * 2;              // 第二轮操作产生的元素个数
        int iter3AlignStart = RoundUp(iter2OutputCount, elementsPerBlock) * elementsPerBlock; // 第三轮操作起始位置偏移，即第二轮产生的元素个数按照datablock(32字节)向上对齐的结果
        int iter3OutputCount = RoundUp(iter2OutputCount, elementsPerRepeat) * 2;              // 第三轮操作产生的元素个数
        int iter3AlignEnd = RoundUp(iter3OutputCount, elementsPerBlock) * elementsPerBlock;   // 第三轮产生的元素个数按照datablock(32字节)向上对齐的结果
        int finalWorkLocalNeedSize = iter2AlignStart + iter3AlignStart + iter3AlignEnd;       // 最终workLocal所需的空间大小  
        pipe.InitBuffer(workQueue, BUFFER_NUM, finalWorkLocalNeedSize * sizeof(float));  
    }
    __aicore__ inline void Process() {

        this->progressUni = 0;
        for(int32_t batch = 0; batch < this->num_batches; batch++) {
            for(int32_t classm = 0; classm < this->num_classes; classm++) {
                CopyInScores(batch * this->num_classes + classm);
                CopyInBoxes(batch);
                CheckIOU(batch, classm);
                // for(int32_t selectedbox = 0; selectedbox < this->max_output_boxes_per_class; selectedbox++) {
                //     int32_t boxindex = CheckIOU(batch, class, selectedbox);
                //     selectedbox = boxindex;
                // }
            }
        }

        //int32_t loopCount = this->num_batches * this->num_classes * this->max_output_boxes_per_class; 
    }

private:
    __aicore__ inline void CopyInScores(int32_t progress) {
        LocalTensor<T> ScoreLocal = inQscores.AllocTensor<T>();
        DataCopy(ScoreLocal, scoresGm[progress * this->spatial_dimension], this->spatial_dimensional);
        inQscores.EnQue(ScoreLocal);        
    }
    __aicore__ inline void CopyInBoxes(int32_t progress) {
        LocalTensor<T> BoxesLocal = inQboxes.AllocTensor<T>();
        DataCopy(BoxesLocal, boxesGm[progress * this->spatial_dimension * 4], this->inQboxeslength);
        inQboxes.EnQue(BoxesLocal);   
    }    
    __aicore__ inline void CheckIOU(int32_t batch, int32_t classm) {
        LocalTensor<T> ScoreLocal = inQscores.DeQue<T>();
        LocalTensor<T> BoxesLocal = inQboxes.DeQue<T>();
        LocalTensor<float> tmpAera1 = tmpBufferA1.Get<float>();
        LocalTensor<float> tmpAera2 = tmpBufferA2.Get<float>();
        LocalTensor<float> tmpAera3 = tmpBufferA3.Get<float>();
        LocalTensor<float> tmpAera4 = tmpBufferA4.Get<float>();
        LocalTensor<float> tmpAera5 = tmpBufferA5.Get<float>();

        LocalTensor<float> workLocal = workQueue.AllocTensor<float>();
        LocalTensor<uint32_t> tmpP1 = tmpBufferP1.Get<uint32_t>();
        LocalTensor<uint32_t> tmpP2 = tmpBufferP2.Get<uint32_t>();
        LocalTensor<uint32_t> tmpP3 = tmpBufferP3.Get<uint32_t>();
        LocalTensor<uint32_t> tmpP4 = tmpBufferP4.Get<uint32_t>();
        LocalTensor<float> tmpX1 = tmpBufferX1.Get<float>();
        LocalTensor<float> tmpX2 = tmpBufferX2.Get<float>();
        LocalTensor<float> tmpY1 = tmpBufferY1.Get<float>();
        LocalTensor<float> tmpY2 = tmpBufferY2.Get<float>();
        LocalTensor<float> SelectedX1 = tmpBufferSX1.Get<float>();
        LocalTensor<float> SelectedX2 = tmpBufferSX2.Get<float>();
        LocalTensor<float> SelectedY1 = tmpBufferSY1.Get<float>();
        LocalTensor<float> SelectedY2 = tmpBufferSY2.Get<float>();
        LocalTensor<uint8_t> cmp1 = QueueTmp2.Get<uint8_t>();

        uint32_t mask = 64;
        uint64_t rsvdCnt = 0;
        // uint32_t patternele1 = 0b00010001000100010001000100010001;
        // uint32_t patternele2 = 0b00100010001000100010001000100010;
        // uint32_t patternele3 = 0b01000100010001000100010001000100;
        // uint32_t patternele4 = 0b10001000100010001000100010001000;
        uint32_t patternele4 = 0b00010001000100010001000100010001;
        uint32_t patternele3 = 0b00100010001000100010001000100010;
        uint32_t patternele2 = 0b01000100010001000100010001000100;
        uint32_t patternele1 = 0b10001000100010001000100010001000;

        tmpP1.SetValue(0, patternele1);
        tmpP1.SetValue(1, patternele1);
        tmpP2.SetValue(0, patternele2);
        tmpP2.SetValue(1, patternele2);
        tmpP3.SetValue(0, patternele3);
        tmpP3.SetValue(1, patternele3);
        tmpP4.SetValue(0, patternele4);
        tmpP4.SetValue(1, patternele4);

        uint16_t repeatTimes = this->spatial_dimension * 4 % 64 ? this->spatial_dimension * 4 / 64 + 1 : this->spatial_dimension * 4 / 64;

        for(int32_t selectedbox = 0; selectedbox < this->max_output_boxes_per_class; selectedbox++) {
            ReduceMax(tmpX1, ScoreLocal, workLocal, this->spatial_dimension, true);
            float scoremax1 = tmpX1.GetValue(0);
            if(scoremax1 >= this->score_threshold) {
                float tmpvalue1 = tmpX1.GetValue(1);
                uint32_t* ptr1 = reinterpret_cast<uint32_t*>(&tmpvalue1); // 获取 half 的位模式
                uint32_t realIndex1 = *reinterpret_cast<int*>(ptr1);
                //ScoreLocal.SetValue(realIndex1, float(-1));
                selected_indicesGm.SetValue(this->progressUni * 3, batch);
                selected_indicesGm.SetValue(this->progressUni * 3 + 1, classm);
                selected_indicesGm.SetValue(this->progressUni * 3 + 2, realIndex1);
                this->progressUni += 1;
                if(this->center_point_box == 0) {
                    float y1 = BoxesLocal.GetValue(realIndex1 * 4);
                    float x1 = BoxesLocal.GetValue(realIndex1 * 4 + 1);
                    float y2 = BoxesLocal.GetValue(realIndex1 * 4 + 2);
                    float x2 = BoxesLocal.GetValue(realIndex1 * 4 + 3);
                    float sx1 = (x1 <= x2) ? x1 : x2;
                    float sx2 = (x1 <= x2) ? x2 : x1;     //  sx2 >= sx1
                    float sy1 = (y1 <= y2) ? y1 : y2;
                    float sy2 = (y1 <= y2) ? y2 : y1;
                    Duplicate(SelectedY1, sy1, this->spatial_dimension);
                    Duplicate(SelectedX1, sx1, this->spatial_dimension);
                    Duplicate(SelectedY2, sy2, this->spatial_dimension);
                    Duplicate(SelectedX2, sx2, this->spatial_dimension);

                    GatherMask(tmpX2, BoxesLocal, tmpP1, true, mask, { 1, repeatTimes, 8, 0 }, rsvdCnt);
                    GatherMask(tmpY2, BoxesLocal, tmpP2, true, mask, { 1, repeatTimes, 8, 0 }, rsvdCnt);
                    GatherMask(tmpX1, BoxesLocal, tmpP3, true, mask, { 1, repeatTimes, 8, 0 }, rsvdCnt);
                    GatherMask(tmpY1, BoxesLocal, tmpP4, true, mask, { 1, repeatTimes, 8, 0 }, rsvdCnt);
                
                    Sub(tmpAera1, SelectedY1, SelectedY2, this->spatial_dimension);
                    Sub(tmpAera2, SelectedX1, SelectedX2, this->spatial_dimension);
                    Mul(tmpAera1, tmpAera1, tmpAera2, this->spatial_dimension);
                    Abs(tmpAera1, tmpAera1, this->spatial_dimension);      //  tmpAera1 is area of Selected box 
                    Sub(tmpAera2, tmpY1, tmpY2, this->spatial_dimension);
                    Sub(tmpAera3, tmpX1, tmpX2, this->spatial_dimension);
                    Mul(tmpAera2, tmpAera2, tmpAera3, this->spatial_dimension);
                    Abs(tmpAera2, tmpAera2, this->spatial_dimension);      //  tmpAera2 is area of input boxes 
                    Max(tmpAera3, tmpX1, tmpX2, this->spatial_dimension);
                    Min(tmpAera3, tmpAera3, SelectedX2, this->spatial_dimension);
                    Min(tmpAera4, tmpX1, tmpX2, this->spatial_dimension);
                    Max(tmpAera4, tmpAera4, SelectedX1, this->spatial_dimension);
                    Sub(tmpAera3, tmpAera3, tmpAera4, this->spatial_dimension);
                    Maxs(tmpAera3, tmpAera3, float(0), this->spatial_dimension);    //tmpAera3 is unite x 

                    Max(tmpAera4, tmpY1, tmpY2, this->spatial_dimension);
                    Min(tmpAera4, tmpAera4, SelectedY2, this->spatial_dimension);
                    Min(tmpAera5, tmpY1, tmpY2, this->spatial_dimension);
                    Max(tmpAera5, tmpAera5, SelectedY1, this->spatial_dimension);
                    Sub(tmpAera4, tmpAera4, tmpAera5, this->spatial_dimension);
                    Maxs(tmpAera4, tmpAera4, float(0), this->spatial_dimension);    //tmpAera4 is unite Y 
                    Mul(tmpAera3, tmpAera3, tmpAera4, this->spatial_dimension);     //tmpAera3 is unite area
                    Add(tmpAera1, tmpAera1, tmpAera2, this->spatial_dimension);
                    Sub(tmpAera1, tmpAera1, tmpAera3, this->spatial_dimension);     //tmpAera1 is merge area
                    Div(tmpAera1, tmpAera3, tmpAera1, this->spatial_dimension);     //tmpAera1 is IOU

                    Duplicate(tmpAera2, float(this->iou_threshold), this->spatial_dimension);
                    Compare(cmp1, tmpAera1, tmpAera2, CMPMODE::GT, this->spatial_dimension); 
                    Duplicate(tmpAera2, float(0), this->spatial_dimension);       
                    Select(tmpAera3, cmp1, tmpAera2, float(1), SELMODE::VSEL_TENSOR_SCALAR_MODE, this->spatial_dimension);
                    Mul(ScoreLocal, ScoreLocal, tmpAera3, this->spatial_dimension);
                } else if (this->center_point_box == 1) {
                    float xc = BoxesLocal.GetValue(realIndex1 * 4);
                    float yc = BoxesLocal.GetValue(realIndex1 * 4 + 1);
                    float width = BoxesLocal.GetValue(realIndex1 * 4 + 2);
                    float height = BoxesLocal.GetValue(realIndex1 * 4 + 3);
                    float sx1 = xc - (width / 2);
                    float sx2 = xc + (width / 2);
                    float sy1 = yc - (height / 2);
                    float sy2 = yc + (height / 2);
                    Duplicate(SelectedY1, sy1, this->spatial_dimension);
                    Duplicate(SelectedX1, sx1, this->spatial_dimension);
                    Duplicate(SelectedY2, sy2, this->spatial_dimension);
                    Duplicate(SelectedX2, sx2, this->spatial_dimension);

                    GatherMask(tmpAera1, BoxesLocal, tmpP1, true, mask, { 1, repeatTimes, 8, 0 }, rsvdCnt);     //height
                    GatherMask(tmpAera2, BoxesLocal, tmpP2, true, mask, { 1, repeatTimes, 8, 0 }, rsvdCnt);     //width
                    GatherMask(tmpAera3, BoxesLocal, tmpP3, true, mask, { 1, repeatTimes, 8, 0 }, rsvdCnt);     //y_center
                    GatherMask(tmpAera4, BoxesLocal, tmpP4, true, mask, { 1, repeatTimes, 8, 0 }, rsvdCnt);    //x_center
                    // Abs(tmpAera1, tmpAera1, this->spatial_dimension);
                    // Abs(tmpAera2, tmpAera2, this->spatial_dimension);
                    Muls(tmpAera1, tmpAera1, float(0.5), this->spatial_dimension);  //height / 2
                    Muls(tmpAera2, tmpAera2, float(0.5), this->spatial_dimension);  //width / 2
                    Add(tmpX2, tmpAera4, tmpAera2, this->spatial_dimension);
                    Sub(tmpX1, tmpAera4, tmpAera2, this->spatial_dimension);
                    Add(tmpY2, tmpAera3, tmpAera1, this->spatial_dimension);
                    Sub(tmpY1, tmpAera3, tmpAera1, this->spatial_dimension);

                    Sub(tmpAera1, SelectedY1, SelectedY2, this->spatial_dimension);
                    Sub(tmpAera2, SelectedX1, SelectedX2, this->spatial_dimension);
                    Mul(tmpAera1, tmpAera1, tmpAera2, this->spatial_dimension);
                    Abs(tmpAera1, tmpAera1, this->spatial_dimension);      //  tmpAera1 is area of Selected box 
                    Sub(tmpAera2, tmpY1, tmpY2, this->spatial_dimension);
                    Sub(tmpAera3, tmpX1, tmpX2, this->spatial_dimension);
                    Mul(tmpAera2, tmpAera2, tmpAera3, this->spatial_dimension);
                    Abs(tmpAera2, tmpAera2, this->spatial_dimension);      //  tmpAera2 is area of input boxes 
                    Max(tmpAera3, tmpX1, tmpX2, this->spatial_dimension);
                    Min(tmpAera3, tmpAera3, SelectedX2, this->spatial_dimension);
                    Min(tmpAera4, tmpX1, tmpX2, this->spatial_dimension);
                    Max(tmpAera4, tmpAera4, SelectedX1, this->spatial_dimension);
                    Sub(tmpAera3, tmpAera3, tmpAera4, this->spatial_dimension);
                    Maxs(tmpAera3, tmpAera3, float(0), this->spatial_dimension);    //tmpAera3 is unite x 

                    Max(tmpAera4, tmpY1, tmpY2, this->spatial_dimension);
                    Min(tmpAera4, tmpAera4, SelectedY2, this->spatial_dimension);
                    Min(tmpAera5, tmpY1, tmpY2, this->spatial_dimension);
                    Max(tmpAera5, tmpAera5, SelectedY1, this->spatial_dimension);
                    Sub(tmpAera4, tmpAera4, tmpAera5, this->spatial_dimension);
                    Maxs(tmpAera4, tmpAera4, float(0), this->spatial_dimension);    //tmpAera4 is unite Y 
                    Mul(tmpAera3, tmpAera3, tmpAera4, this->spatial_dimension);     //tmpAera3 is unite area
                    Add(tmpAera1, tmpAera1, tmpAera2, this->spatial_dimension);
                    Sub(tmpAera1, tmpAera1, tmpAera3, this->spatial_dimension);     //tmpAera1 is merge area
                    Div(tmpAera1, tmpAera3, tmpAera1, this->spatial_dimension);     //tmpAera1 is IOU

                    Duplicate(tmpAera2, float(this->iou_threshold), this->spatial_dimension);
                    Compare(cmp1, tmpAera1, tmpAera2, CMPMODE::GT, this->spatial_dimension); 
                    Duplicate(tmpAera2, float(0), this->spatial_dimension);       
                    Select(tmpAera3, cmp1, tmpAera2, float(1), SELMODE::VSEL_TENSOR_SCALAR_MODE, this->spatial_dimension);
                    Mul(ScoreLocal, ScoreLocal, tmpAera3, this->spatial_dimension);
                }

            } else {
                selectedbox = this->max_output_boxes_per_class;
            }
        }
        inQscores.FreeTensor(ScoreLocal);
        inQboxes.FreeTensor(BoxesLocal);
        workQueue.FreeTensor(workLocal);

    }

    __aicore__ inline int RoundUp(int a, int b){ 
    	return (a + b - 1) / b;
    }

private:
    TPipe pipe;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQboxes, inQscores;
    //TQue<QuePosition::VECOUT, BUFFER_NUM> outQindices;
    TQue<QuePosition::VECOUT, 1> workQueue;
    TBuf<QuePosition::VECCALC> tmpBufferP1, tmpBufferP2, tmpBufferP3, tmpBufferP4;
    TBuf<QuePosition::VECCALC> tmpBufferX1, tmpBufferX2, tmpBufferY1, tmpBufferY2;
    TBuf<QuePosition::VECCALC> tmpBufferSX1, tmpBufferSX2, tmpBufferSY1, tmpBufferSY2;
    TBuf<QuePosition::VECCALC> tmpBufferA1, tmpBufferA2, tmpBufferA3, tmpBufferA4, tmpBufferA5;
    TBuf<QuePosition::VECCALC> QueueTmp2;

    GlobalTensor<T> boxesGm;
    GlobalTensor<T> scoresGm;
    GlobalTensor<int32_t> selected_indicesGm;
    int32_t center_point_box;
    int32_t num_batches;
    int32_t num_classes;
    int32_t spatial_dimension;
    int32_t max_output_boxes_per_class;
    float iou_threshold;
    float score_threshold;
    uint32_t tileDataNum;
    uint32_t inQboxeslength;
    uint32_t spatial_dimensional;
    int32_t progressUni;
};


extern "C" __global__ __aicore__ void non_max_suppression(GM_ADDR boxes, GM_ADDR scores, GM_ADDR max_output_boxes_per_class, GM_ADDR iou_threshold, GM_ADDR score_threshold, GM_ADDR selected_indices, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tdata, tiling);
    KernelNonMaxSuppression<float> op;
    op.Init(boxes, scores, selected_indices, 
            tdata.center_point_box, tdata.num_batches, tdata.num_classes, tdata.spatial_dimension, tdata.max_output_boxes_per_class,
            tdata.iou_threshold, tdata.score_threshold, tdata.tileDataNum);  
    op.Process();
}